#ifndef MYAMG_JACOBISMOOTHER_CUH
#define MYAMG_JACOBISMOOTHER_CUH

#include <CSRSmoother.cuh>
#include <memory>
#include <iostream>
#include <utils/cudaProfiler.cuh>

namespace myamg {
    template<typename Integer, typename Real, typename Allocator>
    class JacobiSmoother : public CSRSmoother<Integer, Real, Allocator, JacobiSmoother<Integer, Real, Allocator> > {
    public:
        using matrixType = typename CSRSmoother<Integer, Real, Allocator, JacobiSmoother>::matrixType;
        using vectorType = typename CSRSmoother<Integer, Real, Allocator, JacobiSmoother>::vectorType;
        using integerType = typename CSRSmoother<Integer, Real, Allocator, JacobiSmoother>::integerType;
        using realType = typename CSRSmoother<Integer, Real, Allocator, JacobiSmoother>::realType;
        using allocatorType = typename CSRSmoother<Integer, Real, Allocator, JacobiSmoother>::allocatorType;

        explicit JacobiSmoother(const double relaxationFactor) : CSRSmoother<Integer, Real, Allocator,
                                                                     JacobiSmoother>(),
                                                                 relaxationFactor(relaxationFactor),
                                                                 recipD(nullptr), z(nullptr) {
        }

        void setupImpl(const matrixType &A, const bool reuse = false) {
            const auto numRows = A.numRows();
            if (!reuse) {
                recipD = std::make_unique<vectorType>(numRows);
                z = std::make_unique<vectorType>(numRows);
            }
            if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                calcRecipDHost(recipD->values(), A.rowPtr(), A.colIndex(), A.values(), numRows);
            } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                dim3 block(512);
                dim3 grid((numRows + block.x - 1) / block.x);
                calcRecipDDevice<<<grid, block>>>(recipD->values(), A.rowPtr(), A.colIndex(), A.values(), numRows);
            } else {
                std::cerr << "JacobiSmoother::setupImpl: Unsupported allocator type" << std::endl;
            }
        }

        void applyImpl(const matrixType &A, const vectorType &b, vectorType &x, const int step,
                       const bool x0Zero = false) {
            const auto numRows = A.numRows();
            double omega = relaxationFactor;

            for (auto k = 0; k < step; ++k) {
                if (!x0Zero || k > 0) {
                    if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                        jacobiUpdateXandZHost(x.values(), z->values(), recipD->values(), A.rowPtr(), A.colIndex(),
                                              A.values(), b.values(), numRows, omega);
                    } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                        dim3 block(512);
                        dim3 grid((numRows + block.x - 1) / block.x);
                        jacobiUpdateZDevice<<<grid, block>>>(z->values(), x.values(), A.rowPtr(), A.colIndex(),
                                                             A.values(), b.values(), numRows);
                        jacobiUpdateXDevice<<<grid, block>>>(x.values(), z->values(), recipD->values(), numRows, omega);
                    } else {
                        std::cerr << "JacobiSmoother::applyImpl: Unsupported allocator type" << std::endl;
                    }
                } else {
                    if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                        jacobiUpdateXZeroHost(x.values(), recipD->values(), b.values(), numRows, omega);
                    } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                        dim3 block(512);
                        dim3 grid((numRows + block.x - 1) / block.x);
                        jacobiUpdateXZeroDevice<<<grid, block>>>(x.values(), recipD->values(), b.values(), numRows,
                                                                 omega);
                    } else {
                        std::cerr << "JacobiSmoother::applyImpl: Unsupported allocator type" << std::endl;
                    }
                }
            }
        }

        void preImpl(const matrixType &A, const vectorType &b, vectorType &x, const int step,
                     const bool x0Zero = false) {
            applyImpl(A, b, x, step, x0Zero);
        }

        void postImpl(const matrixType &A, const vectorType &b, vectorType &x, const int step) {
            applyImpl(A, b, x, step, false);
        }

        void preResidImpl(const matrixType &A, const vectorType &b, vectorType &x, vectorType &r,
                          const int step, const bool x0Zero = false) {
            preImpl(A, b, x, step, x0Zero);
            // r = b - A * x
            if constexpr (AllocatorType<allocatorType>::value == MemoryType::Host) {
                calcResidHost(r.values(), A.rowPtr(), A.colIndex(), A.values(), x.values(), b.values(), A.numRows());
            } else if constexpr (AllocatorType<allocatorType>::value == MemoryType::Cuda) {
                dim3 block(512);
                dim3 grid((A.numRows() + block.x - 1) / block.x);
                calcResidDevice<<<grid, block>>>(r.values(), A.rowPtr(), A.colIndex(), A.values(), x.values(),
                                                 b.values(), A.numRows());
            } else {
                std::cerr << "JacobiSmoother::preResidImpl: Unsupported allocator type" << std::endl;
            }
        }

        double relaxationFactor;

    private:
        std::unique_ptr<vectorType> recipD;
        std::unique_ptr<vectorType> z;
    };
}

#endif //MYAMG_JACOBISMOOTHER_CUH
